# -*- coding:utf-8 -*-
import json
import collections
from trainer.predictor import Predictor


class Inference(object):
    def __init__(self, config_path):
        self._config = json.load(open(config_path, 'r', encoding='utf-8'))
        predictors = collections.OrderedDict()

        for model in self._config:
            predictor = Predictor(lib_dir=self._config[model]['lib_dir'],
                                  ckpt_file=self._config[model]['ckpt_path'])
            predictors[model] = predictor

        self.predictors = predictors
        pass

    def infer(self, query_words, candidate_words):

        ret_dict = collections.OrderedDict()
        import time
        for model in self.predictors:
            start = time.time()
            score = self.predictors[model].predict(query_words, candidate_words)
            ret_dict[model] = score
            # print("{}:\t{}, time cost:\t{:.4f}s".format(model, score, time.time() - start))
        return ret_dict
        pass


def predict_stc(input_path, output_path):
    inference = Inference(config_path='./configs/stc.infer.config.json')
    with open(input_path, 'r', encoding='utf-8') as fr, \
            open(output_path, 'w', encoding='utf-8') as fw:
        for line in fr:
            ss = line.rstrip().split('\t')
            assert len(ss) == 3

            label, query, candidate = ss[0], ''.join(ss[1].split()), ''.join(ss[2].split())
            score = inference.infer(query, candidate)
            score_str = "\t".join(["{:.8f}".format(score[key]) for key in score])
            fw.write("{}\t{}\t{}\t{}\n".format(label, query, candidate, score_str))


if __name__ == '__main__':
    inference = Inference(config_path='./configs/stc.infer.config.json')
    inference.infer("二战结束后，苏军士兵在德国街头强抢自行车，旁边的德国人在围观。", "二战之前有这么耀武扬威的时刻吗")
    # predict_stc("./data/stc/stc.train.tokens.txt", "./data/stc/stc.train.res.txt")
    pass
